import copy
import math

import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
import random
import time

# print('请输入网络的顶点个数：')
# NETWORK_SIZE = int(input())
# print('请输入连边概率：')
# PROBABILITY_OF_EAGE = float(input())
NETWORK_SIZE = 1000
PROBABILITY_OF_EAGE = 0.006
M = 3
S0 = 0.9 # ratio
I0 = 0.1 # ratio
R0 = 0.0 # ratio
beta = 0.2
gamma = 0.5
T = 50
# er = np.zeros((NETWORK_SIZE, NETWORK_SIZE), dtype=int)  # 初始化邻接矩阵
# ba = np.zeros((NETWORK_SIZE, NETWORK_SIZE), dtype=int)  # 初始化邻接矩阵
random.seed(time.time())  # 'random.random()#生成[0,1)之间的随机数



# 生成ER网络
def get_er():
    # print('请输入网络的顶点个数：')
    # NETWORK_SIZE = int(input())
    # print('请输入连边概率：')
    # PROBABILITY_OF_EAGE = float(input())
    # er = np.zeros((NETWORK_SIZE, NETWORK_SIZE), dtype=int)  # 初始化邻接矩阵
    # count = 0
    # probability = 0.0
    # for i in range(NETWORK_SIZE):
    #     for j in range(i + 1, NETWORK_SIZE):
    #         probability = random.random()
    #         if probability < PROBABILITY_OF_EAGE:
    #             count = count + 1
    #             er[i][j] = er[j][i] = 1
    # print(er)

    # 使用networkx框架
    ERNX = nx.erdos_renyi_graph(NETWORK_SIZE,PROBABILITY_OF_EAGE)

    er = np.array(nx.adjacency_matrix(ERNX,dtype=int).todense())

    neighbors = [[]]*NETWORK_SIZE

    status = [0] * NETWORK_SIZE

    for i in range(0,NETWORK_SIZE):
        neighbors[i] = (list(ERNX.neighbors(i)))

    # print(er)
    return er,neighbors,status

# 生成BA网络
def get_ba():
    # 使用networkx框架
    BANX=nx.barabasi_albert_graph(NETWORK_SIZE,M)

    ba = np.array(nx.adjacency_matrix(BANX,dtype=int).todense())

    neighbors = [[]]*NETWORK_SIZE

    status = [0] * NETWORK_SIZE

    for i in range(0,NETWORK_SIZE):
        neighbors[i] = (list(BANX.neighbors(i)))

    # print(ba)
    return ba,neighbors,status

# 显示网络图
def showGraph(matrix):
    G = nx.Graph()
    for i in range(len(matrix)):
        for j in range(len(matrix)):
            if matrix[i][j] == 1:  # 如果不加这句将生成完全图，ER网络的邻接矩阵将不其作用
                G.add_edge(i, j)
    nx.draw(G,with_labels=True)
    plt.show()

# 打印度分布图
def print_degree(matrix):

    row_num = len(matrix)
    degree = [0] * row_num
    for i in range(0,row_num):
        degree[i] = sum(matrix[i])
    print(degree)
    same_num = get_same_num(degree)
    plt.plot(range(0,row_num),same_num,marker='o', mec='r')
    plt.xlabel("degree")  # X轴标签
    plt.ylabel("P(degree)")  # Y轴标签
    plt.title("Degree distribution diagram")  # 标题
    plt.show()

# 获取度相同的个数
def get_same_num(degree):
    total_num = len(degree)
    same_num = [0] * total_num
    for i in range(0,total_num):
        same_num[degree[i]] += 1
    for j in range(0,total_num):
        same_num[j] /= total_num
    return same_num

# 计算该节点邻接节点中感染数
def get_i_num(status,neighbors):
    inum = 0
    for neighbor in neighbors:
        if status[neighbor] == 1:
            inum += 1
    return inum

# 计算该节点是否被感染
def is_to_i(inum):
    to_i_probablity = 1 - math.pow((1 - beta), inum)
    probability = random.random()
    if probability<=to_i_probablity:
        return True
    else:
        return False

# 计算该节点是否康复
def is_to_r():
    probability = random.random()
    if probability<=gamma:
        return True
    else:
        return False

def init_i(status,I0):

    for i in range(0,int(NETWORK_SIZE*I0)):
        j = round(random.random()*NETWORK_SIZE)
        # print(j)
        status[j-1] = 1
    return status


def SIR(status,neighbors,status2,neighbors2):
    snums = [0] * T
    inums = [0] * T
    rnums = [0] * T
    snums2 = [0] * T
    inums2 = [0] * T
    rnums2 = [0] * T
    totals = [0] * T
    totali = [0] * T
    totalr = [0] * T
    inums[0] = status.count(1)
    snums[0] = status.count(0)
    inums2[0] = status2.count(1)
    snums2[0] = status2.count(0)
    totali[0] = status2.count(1) + status.count(1)
    totals[0] = status2.count(0) + status.count(0)

    for t in range(1,T):
        for i in range(0,NETWORK_SIZE):
            #网络1判断
            if status[i] == 0: #该节点为s
                inum = get_i_num(status,neighbors[i])
                if is_to_i(inum):
                    status[i] = 1 #该节点被感染
            elif status[i] == 1: #该节点为i
                if is_to_r():
                    status[i] = 2 #该节点康复
            #网络2判断
            if status2[i] == 0: #该节点为s
                inum2 = get_i_num(status2,neighbors2[i])
                if is_to_i(inum2):
                    status2[i] = 1 #该节点被感染
            elif status2[i] == 1: #该节点为i
                if is_to_r():
                    status2[i] = 2 #该节点康复
        for i in range(0, NETWORK_SIZE):
            if status[i] == 0: #该节点为s
                if status2[i] == 1:
                    status[i] = 1
            elif status[i] == 1:
                if status2[i] == 0:
                    status2[i] = 1

        snum = status.count(0)
        inum = status.count(1)
        rnum = status.count(2)
        snum2 = status2.count(0)
        inum2 = status2.count(1)
        rnum2 = status2.count(2)
        totals[t] = snum + snum2
        totali[t] = inum + inum2
        totalr[t] = rnum + rnum2

        # snums[t] = snum
        # inums[t] = inum
        # rnums[t] = rnum
        # snums2[t] = snum2
        # inums2[t] = inum2
        # rnums2[t] = rnum2

    return totals,totali,totalr

def print_SIR(snums,inums,rnums):
    print(snums)
    print(inums)
    print(rnums)

    plt.title('SIR model results')
    plt.xlabel('t')
    plt.ylabel('N(t)')
    plt.plot(range(0,T),snums)
    plt.plot(range(0,T),inums)
    plt.plot(range(0,T),rnums)
    plt.legend(['S','I','R'])
    plt.show()


if __name__ == '__main__':

    #获取er表，每个结点的邻接节点表，每个结点的状态表
    er,neighbors,status = get_er()
    #初始化节点状态
    status = init_i(status,I0)
    # 获取ba表，每个结点的邻接节点表，每个结点的状态表
    ba,neighbors2,status2 = get_ba()
    #初始化节点状态
    status2 = init_i(status2,I0)

    # print_degree(er)
    # showGraph(er)
    # print_degree(ba)
    # showGraph(ba)

    # totals,totali,totalr = SIR(status,neighbors,status2,neighbors2)
    # print_SIR(totals,totali,totalr)


    totals = [0] * T
    totali = [0] * T
    totalr = [0] * T
    for i in range(0,100):
        s1 = copy.deepcopy(status)
        s2 = copy.deepcopy(status2)
        n1 = copy.deepcopy(neighbors)
        n2 = copy.deepcopy(neighbors2)
        snums,inums,rnums = SIR(s1,n1,s2,n2)
        totals += np.array(snums)
        totali += np.array(inums)
        totalr += np.array(rnums)

    print_SIR(totals,totali,totalr)


